import os
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io

# 初始化 InferenceClient
def init_client():
    """初始化 HuggingFace InferenceClient"""
    try:
        client = InferenceClient(
            api_key=os.environ.get("HF_API_KEY", ""),  # 可选，免费模型可以不设置
        )
        return client
    except Exception as e:
        raise gr.Error(f"客户端初始化失败: {str(e)}")

def compress_image(image, max_size_mb=3.5, quality=85):
    """
    压缩图像以确保不超过最大文件大小
    """
    max_size_bytes = max_size_mb * 1024 * 1024
    
    # 首先尝试调整图像尺寸
    width, height = image.size
    scale_factor = 1.0
    
    # 如果图像太大，按比例缩小
    if width > 1024 or height > 1024:
        scale_factor = min(1024 / width, 1024 / height)
        new_width = int(width * scale_factor)
        new_height = int(height * scale_factor)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 转换为 RGB 模式（如果不是的话）
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # 尝试不同的质量设置
    for q in range(quality, 30, -10):
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='JPEG', quality=q, optimize=True)
        img_byte_arr.seek(0)
        
        if img_byte_arr.tell() <= max_size_bytes:
            img_byte_arr.seek(0)
            return Image.open(img_byte_arr)
    
    # 如果还是太大，进一步缩小尺寸
    while scale_factor > 0.3:
        scale_factor *= 0.8
        new_width = int(width * scale_factor)
        new_height = int(height * scale_factor)
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        img_byte_arr = io.BytesIO()
        resized_image.save(img_byte_arr, format='JPEG', quality=60, optimize=True)
        img_byte_arr.seek(0)
        
        if img_byte_arr.tell() <= max_size_bytes:
            img_byte_arr.seek(0)
            return Image.open(img_byte_arr)
    
    raise gr.Error("无法将图像压缩到合适的大小，请尝试使用更小的图像")

def process_image_to_image(input_image, prompt, model_name):
    """
    使用指定的模型进行图像到图像转换
    """
    try:
        # 初始化客户端
        client = init_client()
        
        # 压缩图像
        compressed_image = compress_image(input_image)
        
        # 将 PIL Image 转换为字节
        img_byte_arr = io.BytesIO()
        compressed_image.save(img_byte_arr, format='JPEG', quality=85, optimize=True)
        img_byte_arr = img_byte_arr.getvalue()
        
        # 检查文件大小
        file_size_mb = len(img_byte_arr) / (1024 * 1024)
        print(f"图像文件大小: {file_size_mb:.2f} MB")
        
        if file_size_mb > 4:
            raise gr.Error(f"图像文件仍然过大 ({file_size_mb:.2f} MB)，请使用更小的图像")
        
        # 调用模型进行图像转换
        output_image = client.image_to_image(
            img_byte_arr,
            prompt=prompt,
            model=model_name,
        )
        
        return output_image
        
    except Exception as e:
        if "402" in str(e) or "Payment Required" in str(e):
            raise gr.Error("免费额度已用完，请升级到 PRO 计划或使用其他免费模型")
        elif "413" in str(e) or "Payload Too Large" in str(e):
            raise gr.Error("图像文件过大，请使用更小的图像或尝试不同的图像")
        else:
            raise gr.Error(f"图像处理失败: {str(e)}")

def generate_image(img, text, model):
    """图像生成函数"""
    if img is None:
        raise gr.Error("请先上传一张图片")
    if not text.strip():
        raise gr.Error("请输入提示词")
    
    try:
        result = process_image_to_image(img, text, model)
        return result, "图像生成完成！"
    except Exception as e:
        raise gr.Error(f"处理失败: {str(e)}")

# 创建界面
with gr.Blocks(title="免费图像转换工具", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 免费图像转换工具
    
    使用免费的 AI 模型将您的图像转换为新的风格。上传一张图片，输入提示词，然后生成全新的图像！
    
    ### 使用说明：
    1. 上传一张图片（建议小于 4MB）
    2. 选择免费的模型
    3. 输入描述性的提示词
    4. 点击"生成图像"按钮
    5. 等待处理完成，查看结果
    
    **注意：** 如果图像过大，系统会自动压缩图像以确保处理成功。
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # 输入区域
            gr.Markdown("### 📤 输入")
            
            input_image = gr.Image(
                label="上传图片（建议小于 4MB）",
                type="pil",
                height=300
            )
            
            model_selector = gr.Dropdown(
                choices=[
                    "stabilityai/stable-diffusion-2-1",
                    "runwayml/stable-diffusion-v1-5",
                    "CompVis/stable-diffusion-v1-4"
                ],
                value="stabilityai/stable-diffusion-2-1",
                label="选择免费模型",
                interactive=True
            )
            
            prompt = gr.Textbox(
                label="提示词",
                placeholder="例如：Turn the cat into a tiger, Make it more colorful...",
                lines=3,
                max_lines=5
            )
            
            generate_btn = gr.Button(
                "🎨 生成图像",
                variant="primary",
                size="lg"
            )
            
            # 状态显示
            status = gr.Textbox(
                label="状态",
                value="准备就绪",
                interactive=False
            )
        
        with gr.Column(scale=1):
            # 输出区域
            gr.Markdown("### 📥 输出")
            
            output_image = gr.Image(
                label="生成结果",
                height=400
            )
    
    # 示例区域
    with gr.Accordion("💡 示例提示词", open=False):
        gr.Markdown("""
        **动物转换：**
        - Turn the cat into a tiger
        - Transform the dog into a wolf
        - Make the bird look like an eagle
        
        **风格转换：**
        - Make it more colorful and vibrant
        - Convert to watercolor painting style
        - Transform into anime style
        
        **环境转换：**
        - Put the subject in a forest
        - Place it in a futuristic city
        - Set it in a magical garden
        """)
    
    # 绑定事件
    generate_btn.click(
        fn=generate_image,
        inputs=[input_image, prompt, model_selector],
        outputs=[output_image, status],
        api_name="generate"
    )

if __name__ == "__main__":
    print("启动免费图像转换工具...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7861,
        share=True,
        debug=True
    ) 